#借助ncbi-genome-download工具进行自动化数据下载
import os   
import csv   
import subprocess   
import multiprocessing   
from collections import defaultdict   
import logging
# 定义工作目录   
work_dir = "/home/zhangqiang/zxy/genome_download"
# 配置日志   
logging.basicConfig(
    filename=os.path.join(work_dir, "download_new.log"),
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"   
)
# 检查 ncbi-genome-download 是否可用   
def check_ncbi_genome_download():
    try:
        result = subprocess.run(
            ["which", "ncbi-genome-download"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True
        )
        if result.returncode == 0:
            print("ncbi-genome-download found:", result.stdout.strip())
            logging.info(f"ncbi-genome-download found: {result.stdout.strip()}")
        else:
            print("ncbi-genome-download not found.")
            logging.error(f"ncbi-genome-download not found: {result.stderr}")
            print("Error message:", result.stderr)
    except Exception as e:
        print("An error occurred:", str(e))
        logging.error(f"Error checking ncbi-genome-download: {str(e)}")
check_ncbi_genome_download()
# 读取CSV文件   
csv_file_path = os.path.join(work_dir, "download_new.csv")   
with open(csv_file_path, mode='r', encoding='utf-8') as file:
    reader = csv.DictReader(file)
    data = [row for row in reader]
# 按目（order）分组，保持原始顺序   
orders = defaultdict(list)   
for row in data:
    orders[row['order']].append(row)
# 确保文件夹名称有效   
def sanitize_folder_name(folder_name):
    valid_chars = "".join([char for char in folder_name if char.isalnum() or char in "_-."])  # 保留字母、数字、下划线和点
    return valid_chars.replace(" ", "_")  # 将空格替换为下划线
# 定义下载函数   
def download_genome(accession, group, output_folder):
    logging.info(f"Attempting to download {accession} for group {group}")
    if accession.startswith("GCF_"):
        section = "refseq"
    else:
        section = "genbank"
    cmd = f"ncbi-genome-download --formats fasta,gff,protein-fasta --section {section} --assembly-accessions {accession} --output-folder {output_folder} {group}"
    logging.info(f"Download command: {cmd}")
    try:
        os.makedirs(output_folder, exist_ok=True)
        result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if result.returncode == 0:
            logging.info(f"Download successful for {accession}")
            print(f"Download successful for {accession}")
            return True, result.stdout.decode(), result.stderr.decode()
        else:
            logging.error(f"Download failed for {accession}: {result.stderr.decode()}")
            print(f"Download failed for {accession}")
            return False, result.stderr.decode(), None
    except Exception as e:
        logging.error(f"Error occurred while downloading {accession}: {str(e)}")
        print(f"An error occurred: {str(e)}")
        return False, str(e), None
# 处理每个目      
def process_order(order, rows, processed_accessions):
    print(f"Processing order: {order}")
    logging.info(f"Processing order: {order}")
    found = False
    for row in rows:
        accession = row['assembly_accession']
        if accession in processed_accessions:
            logging.info(f"Skipping duplicate accession: {accession}")
            print(f"Skipping duplicate accession: {accession}")
            continue
        processed_accessions.add(accession)
        group = row['group']
        safe_group = sanitize_folder_name(group)
        output_folder = os.path.join(work_dir, safe_group)
        os.makedirs(output_folder, exist_ok=True)
        # 检查基因组文件是否已存在
        if check_files(output_folder, accession):
            logging.info(f"Files for {accession} already exist in {output_folder}, skipping download.")
            print(f"Files for {accession} already exist in {output_folder}, skipping download.")
            found = True
            break
        success, stdout, stderr = download_genome(accession, group, output_folder)
        if success:
            if check_files(output_folder, accession):
                logging.info(f"Downloaded suitable files for {accession}")
                print(f"Downloaded suitable files for {accession}")
                found = True
                break
            else:
                logging.warning(f"No suitable files found for {accession}, but download was successful.")
        else:
            logging.error(f"Download failed for {accession}: {stderr}")
            print(f"Download failed for {accession}: {stderr}")
            continue  # Skip to the next accession
    if not found:
        logging.info(f"No suitable data found for order: {order}")
        print(f"No suitable data found for order: {order}")
    return f"Order: {order}, Found: {found}"
# 检查文件是否存在   
def check_files(output_folder, accession):
    # 遍历输出文件夹中的所有子文件夹
    for root, dirs, files in os.walk(output_folder):
        for file in files:
# 检查文件名是否包含 accession 并且以 .fna, .faa, .gff 或其压缩格式结尾
            if accession in file and file.endswith(('.fna', '.faa', '.gff', '.fna.gz', '.faa.gz', '.gff.gz')):
                return True
    return False

if __name__ == "__main__":
    # 维护一个全局的已处理 accession 集合
    processed_accessions = set()
    # 按表格顺序处理每个目
    tasks = [(order, rows, processed_accessions) for order, rows in orders.items()]
    with multiprocessing.Pool(processes=8) as pool:
        results = pool.starmap(process_order, tasks)
    # 保存结果
    with open(os.path.join(work_dir, "download_results.txt"), "w") as result_file:
        for result in results:
            result_file.write(result + "\n")
    print("Download process completed.")
    logging.info("Download process completed.")